from epics import PV
import json
from datetime import datetime
from config_loader import get_cavity_configs, get_cavity_names
import os
import time

class EPICSSnapshot:
    def __init__(self):
        # 按腔体初始化所有PV连接
        self.cavity_pvs = {}
        for cavity, pv_dict in get_cavity_configs().items():
            self.cavity_pvs[cavity] = {
                name: PV(pv_name) for name, pv_name in pv_dict.items()
            }

        # 创建保存快照的目录
        self.snapshot_dir = "snapshots"
        if not os.path.exists(self.snapshot_dir):
            os.makedirs(self.snapshot_dir)
        
        # 等待所有PV连接建立
        self._wait_for_connections()

    def _wait_for_connections(self, timeout=5.0):
        """等待所有PV连接建立"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            unconnected = []
            for cavity, pvs in self.cavity_pvs.items():
                for name, pv in pvs.items():
                    if not pv.connected:
                        unconnected.append(f"{cavity}-{name}")
            if not unconnected:
                return True
            time.sleep(0.1)
        
        if unconnected:
            raise ConnectionError(f"无法连接到以下PV: {', '.join(unconnected)}")
        return False

    def _check_pv_access(self, cavity=None):
        """检查指定腔体或所有腔体的PV连接状态"""
        disconnected = []
        if cavity:
            pvs = self.cavity_pvs.get(cavity, {})
            for name, pv in pvs.items():
                if not pv.connected:
                    disconnected.append(f"{cavity}-{name}")
        else:
            for cavity, pvs in self.cavity_pvs.items():
                for name, pv in pvs.items():
                    if not pv.connected:
                        disconnected.append(f"{cavity}-{name}")
        
        if disconnected:
            raise ConnectionError(f"以下PV未连接: {', '.join(disconnected)}")

    def take_snapshot(self, cavity=None, description=""):
        """
        获取指定腔体或所有腔体的PV当前值并保存为快照
        cavity: 腔体名称，如果为None则保存所有腔体
        """
        try:
            self._check_pv_access(cavity)
            
            values = {}
            read_errors = []
            
            if cavity:
                # 只保存指定腔体的值
                cavities_to_save = {cavity: self.cavity_pvs[cavity]}
            else:
                # 保存所有腔体的值
                cavities_to_save = self.cavity_pvs

            for cav_name, pvs in cavities_to_save.items():
                values[cav_name] = {}
                for name, pv in pvs.items():
                    value = pv.get(timeout=1.0)
                    if value is None:
                        read_errors.append(f"{cav_name}-{name}")
                    values[cav_name][name] = value
            
            if read_errors:
                raise RuntimeError(f"无法读取以下PV的值: {', '.join(read_errors)}")

            snapshot_data = {
                'timestamp': datetime.now().strftime('%Y-%m-%d_%H-%M-%S'),
                'description': description,
                'cavity': cavity if cavity else 'all',
                'values': values
            }

            # 生成文件名
            cavity_prefix = f"{cavity}_" if cavity else "all_"
            filename = f"snapshot_{cavity_prefix}{snapshot_data['timestamp']}.json"
            filepath = os.path.join(self.snapshot_dir, filename)

            with open(filepath, 'w') as f:
                json.dump(snapshot_data, f, indent=4)

            return True, filepath

        except Exception as e:
            return False, str(e)

    def load_snapshot(self, filepath, cavity=None):
        """
        从文件加载快照并设置PV值
        cavity: 指定要加载的腔体，如果为None则按照快照文件中的设置加载
        """
        try:
            with open(filepath, 'r') as f:
                snapshot_data = json.load(f)

            snapshot_cavity = snapshot_data.get('cavity', 'all')
            if cavity and snapshot_cavity != 'all' and cavity != snapshot_cavity:
                raise ValueError(f"快照文件仅包含腔体 {snapshot_cavity} 的数据")

            write_errors = []
            values = snapshot_data['values']

            if cavity:
                # 只加载指定腔体的值
                if cavity not in values:
                    raise ValueError(f"快照中不包含腔体 {cavity} 的数据")
                cavities_to_load = {cavity: values[cavity]}
            else:
                # 加载所有腔体的值
                cavities_to_load = values

            self._check_pv_access(cavity)

            for cav_name, cav_values in cavities_to_load.items():
                for name, value in cav_values.items():
                    if name in self.cavity_pvs[cav_name]:
                        success = self.cavity_pvs[cav_name][name].put(value, wait=True, timeout=1.0)
                        if not success:
                            write_errors.append(f"{cav_name}-{name}")
                    else:
                        write_errors.append(f"{cav_name}-{name}(不存在)")

            if write_errors:
                raise RuntimeError(f"无法写入以下PV的值: {', '.join(write_errors)}")

            loaded_cavities = "所有腔体" if not cavity else f"腔体 {cavity}"
            return True, f"已成功加载{loaded_cavities}的快照"

        except Exception as e:
            return False, str(e)

    def list_snapshots(self, cavity=None):
        """列出所有可用的快照"""
        snapshots = []
        for filename in os.listdir(self.snapshot_dir):
            if filename.endswith('.json'):
                filepath = os.path.join(self.snapshot_dir, filename)
                try:
                    with open(filepath, 'r') as f:
                        data = json.load(f)
                        snapshot_cavity = data.get('cavity', 'all')
                        if cavity is None or snapshot_cavity == cavity or snapshot_cavity == 'all':
                            snapshots.append({
                                'filename': filename,
                                'timestamp': data['timestamp'],
                                'description': data['description'],
                                'cavity': snapshot_cavity
                            })
                except Exception as e:
                    print(f"读取快照文件 {filename} 时出错: {str(e)}")
        return snapshots

# 使用示例
if __name__ == "__main__":
    snapshot_manager = EPICSSnapshot()

    while True:
        print("\nEPICS快照管理工具")
        print("1. 创建新快照")
        print("2. 加载快照")
        print("3. 列出所有快照")
        print("4. 退出")

        choice = input("请选择操作 (1-4): ")

        if choice == '1':
            cavity = input("请输入腔体名称 (可选): ")
            description = input("请输入快照描述 (可选): ")
            snapshot_manager.take_snapshot(cavity, description)

        elif choice == '2':
            snapshots = snapshot_manager.list_snapshots()
            if not snapshots:
                print("没有找到可用的快照")
                continue

            print("\n可用的快照:")
            for i, snapshot in enumerate(snapshots, 1):
                print(f"{i}. {snapshot['filename']} - {snapshot['description']}")

            try:
                index = int(input("请选择要加载的快照编号: ")) - 1
                if 0 <= index < len(snapshots):
                    filepath = os.path.join(snapshot_manager.snapshot_dir, snapshots[index]['filename'])
                    snapshot_manager.load_snapshot(filepath)
                else:
                    print("无效的选择")
            except ValueError:
                print("请输入有效的数字")

        elif choice == '3':
            snapshots = snapshot_manager.list_snapshots()
            if not snapshots:
                print("没有找到快照")
            else:
                print("\n可用的快照:")
                for snapshot in snapshots:
                    print(f"文件名: {snapshot['filename']}")
                    print(f"时间戳: {snapshot['timestamp']}")
                    print(f"描述: {snapshot['description']}")
                    print("-" * 50)

        elif choice == '4':
            print("程序退出")
            break

        else:
            print("无效的选择，请重试") 